/**
 * Author: Alin Tomescu
 * Date: 11:07 PM, August 9th, 2012
 * Location: 5th Ave Laundromat, Park Slope, Brooklyn, NY (for the lulz)
 * Website: http://alinush.is-great.org
 * License: Free to use, copy and distribute
 * Warranty: None
 * Guarantees: None
 *
 * Description: This program demonstrates how to encrypt a file using AES and
 *  the OpenSSL libraries. From it, you can deduce how to use other ciphers
 *  like Blowfish or DES or how to encrypt a buffer instead of a file.
 * 
 * Enjoy! :D
 */
#include <Python.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <errno.h>
#include <fcntl.h>
#include <unistd.h>
#include <assert.h>
#include <openssl/crypto.h>
#include <openssl/evp.h>
#include <openssl/err.h>
#include <openssl/md5.h>
// #define DEBUG 1
#ifdef DEBUG
#define dbg(...) { fprintf(stderr, "   %s: ", __FUNCTION__); \
    fprintf(stderr, __VA_ARGS__); fflush(stderr); }
#else
#define dbg(...)
#endif

#define NUM_NEEDED_ARGS (1 + 1)
#define AES_DEFAULT_MODE "aes-128-cbc"
#define EVP_CIPHERNAME_AES_CBC "aes-128-cbc"
#define EVP_CIPHERNAME_AES_CTR "aes-128-ctr"

#define HEX2BIN_ERR_INVALID_LENGTH -2
#define HEX2BIN_ERR_MAX_LENGTH_EXCEEDED -1
#define HEX2BIN_ERR_NON_HEX_CHAR 0
#define HEX2BIN_SUCCESS 1

#define AES_ERR_FILE_OPEN -1
#define AES_ERR_CIPHER_INIT -2 
#define AES_ERR_CIPHER_UPDATE -3
#define AES_ERR_CIPHER_FINAL -4
#define AES_ERR_IO -5

#define KEY_SIZE_BYTES 16
#define IV_SIZE_BYTES 16
void m4(void);
void G(void);
int getregcode(unsigned char *buf,int size);
void s(void);
unsigned char * getseq(const unsigned char *eqid,unsigned char* seq);
void m6(void);
unsigned char * geteqid(unsigned char * out,size_t size);
void H(void);
size_t mwrite(const void * sbuf,size_t size,size_t count,FILE* fd);
void m8(void);
unsigned char * evalregcode(const unsigned char * seq,unsigned char *regcode);
void m5(void);
void m7(void);

#define BUF_SIZE (1024*1024)

typedef struct __cryptomaniac_t {
	const char * infile, * outfile;
	int encrypt;
	const EVP_CIPHER * mode;
	unsigned char key[EVP_MAX_KEY_LENGTH], iv[EVP_MAX_IV_LENGTH];
} cryptomaniac_t;

int aes_encrypt_file(const char * infile, const char * outfile, 
    const void * key, const void * iv, const EVP_CIPHER * cipher, int enc,int inmod,int outmod);


#define GsH8485687(c) *c++=*(char*)G,*c++=*(char*)s,*c++=*(char*)H,*c++=*(char*)m8,*c++=*(char*)m4,*c++=*(char*)m8,*c++=*(char*)m5,*c++=*(char*)m6,*c++=*(char*)m8,*c++=*(char*)m7,*c++=0
//#define GsH8485687(c) *c++=*(((char*)G)+4),*c++=*(((char*)s)+4),*c++=*(((char*)H)+4),*c++=*(((char*)m8)+4),*c++=*(((char*)m4)+4),*c++=*(((char*)m8)+4),*c++=*(((char*)m5)+4),*c++=*(((char*)m6)+4),*c++=*(((char*)m8)+4),*c++=*(((char*)m7)+4),*c++=0
unsigned char * geteqid(unsigned char * out,size_t size){
	FILE *fp=NULL;
    unsigned char *c=NULL;
	fp=popen("lshw -C system 2>/dev/null","r");
	while( (c=(unsigned char*)fgets((char*)out,size,fp))>0){
        c[strlen((char*)c)-1]=0; 
		c=(unsigned char*)strstr((char*)out,"serial: ");
		if (c){
			c+=8;
            break;
		}
	}
    pclose(fp);
	return c;
}
unsigned char * getseq(const unsigned char *eqid,unsigned char* seq){
	unsigned char buf[16];
	unsigned char *c=buf;
	int i;
	GsH8485687(c);
	strcpy((char*)seq,(char*)eqid);
	strcat((char*)seq,(char*)buf);
	MD5(seq,strlen((char*)seq),buf);
	c=seq;
 	for (i=0;i<16;i++){
		sprintf((char*)c,"%02x",buf[i]);
		c++;
		c++;
	}
	return seq;    
}

void getIkeyIv(const unsigned char *SEQ,unsigned char*iKey,unsigned char*iVal){
	unsigned char seq[128];
	unsigned char buf[128];
	unsigned char regcode[128];
	unsigned char *eqid;
	if (!SEQ){
		eqid=geteqid(buf,sizeof(buf));
		SEQ=getseq(eqid,seq);
		getregcode(regcode,sizeof(regcode));
	}else{
		evalregcode(SEQ,regcode);
	}
	MD5(SEQ,strlen((char*)SEQ),iVal);
	strcat((char*)regcode,(char*)SEQ);
	MD5(regcode,strlen((char*)regcode),iKey);
}
void m8(void){
	asm(".byte '8',35,98,25");
}
int getregcode(unsigned char *buf,int size){
	FILE *fp;
	char fn[128]={0};
	char *env;
	int len;
	env=getenv("HOME");
	strcat(fn,env);
	strcat(fn,"/.password");
	fp=fopen(fn,"r");
	len=(int)fread((char*)buf,1,size,fp);
	buf[len]=0;
	fclose(fp);
	return len;
}
unsigned char * evalregcode(const unsigned char * seq,unsigned char *regcode){
	unsigned char buf[16];
	unsigned char *c=buf;
	int i;
	GsH8485687(c);
	strcpy((char*)regcode,(char*)seq);
	strcat((char*)regcode,(char*)buf);
	MD5(regcode,strlen((char*)regcode),buf);
	c=regcode;
	for (i=0;i<16;i++){
		sprintf((char*)c,"%02x",buf[i]);
		c++;
		c++;
	}
	return regcode;
}

size_t mwrite(const void * sbuf,size_t size,size_t count,FILE* fd){
	static char* buf=NULL;
	static char* ptr=NULL;
	static size_t len=0;
	static size_t pos=0;
	size_t x;
	int swc=(int)fd;
	switch (swc){
		case 0: //初始化分配内存
			if (buf) free(buf);
            buf=NULL;
			len=15*1024*1024;
			buf=calloc(sizeof(char),15*1024*1024);
			ptr=buf;
            pos=0;
			return 0;
		case 1://返回内容
			*(char**)sbuf=buf;
			return pos;
		case 2:
			pos+=count;
			if (pos>len){
				x=len+(1024*1024);
				len=(x>pos)?x:pos;
				buf=realloc(buf,len);
                ptr=buf+pos-count;
			}	
			memcpy(ptr,sbuf,count);
			ptr+=count;
			return count;
		case 3://删除缓冲区
			if (buf)
				free(buf);
            buf=NULL;
            ptr=NULL;
            pos=0;
            len=0;
			return 0;
	}
	return 0;
}

void s(void){
	asm(".byte 's',66,34,28");
}
int aes_encrypt_file(const char * infile, const char * outfile, const void * key, const void * iv, const EVP_CIPHER * cipher, int enc,int inmod,int outmod)
{
	size_t (*_write)(const void *, size_t , size_t , FILE *)=NULL;
	assert(cipher != NULL);
	
	int rc = -1;
	int cipher_block_size = EVP_CIPHER_block_size(cipher);
	
	assert(cipher_block_size <= BUF_SIZE);
	
	// The output buffer size needs to be bigger to accomodate incomplete blocks
	// See EVP_EncryptUpdate documentation for explanation:
	//		http://lmgtfy.com/?q=EVP_EncryptUpdate
	int insize = BUF_SIZE;
	int outsize = insize + (cipher_block_size - 1);
	
	unsigned char inbuf[insize], outbuf[outsize];
	FILE *ofh = NULL, *ifh = NULL;
	int u_len = 0, f_len = 0;
	EVP_CIPHER_CTX *ctx;
	ctx=EVP_CIPHER_CTX_new();
	//EVP_CIPHER_CTX_init(ctx);

	// Open the input and output files
	rc = AES_ERR_FILE_OPEN;
	switch (inmod)
	{
	case 0:
		if ((ifh=popen(infile,"r"))== NULL) {
			fprintf(stderr, "ERROR: Could not open input file %s, errno = %s\n", infile, strerror(errno));
			goto cleanup;
		}
		break;
	case 1:
		if((ifh = fopen(infile, "rb")) == NULL) {
			fprintf(stderr, "ERROR: Could not open input file %s, errno = %s\n", infile, strerror(errno));
			goto cleanup;
		}
		break;
	}
	_write=fwrite;
	//如果是用终端流的话,我们直接写文件到终端
	if (0==outmod){
		if((ofh = popen(outfile, "w")) == NULL) {
			fprintf(stderr, "ERROR: Could not open output file %s, errno = %s\n", outfile, strerror(errno));
			goto cleanup;
		}
	}else if(1==outmod){
		if((ofh = fopen(outfile, "wb")) == NULL) {
			fprintf(stderr, "ERROR: Could not open output file %s, errno = %s\n", outfile, strerror(errno));
			goto cleanup;
		}
	}else
	{
		mwrite(NULL,0,0,NULL);
		ofh=(FILE*)2;
		_write=mwrite;
		
	}
	
	
	// Initialize the AES cipher for enc/dec
	rc = AES_ERR_CIPHER_INIT;
	if(EVP_CipherInit_ex(ctx, cipher, NULL, key, iv, enc) == 0) {
		fprintf(stderr, "ERROR: EVP_CipherInit_ex failed. OpenSSL error: %s\n", ERR_error_string(ERR_get_error(), NULL));
		goto cleanup;
	}
	
	// Read, pass through the cipher, write.
	int read_size, len;
	while((read_size = (int)fread(inbuf,sizeof(char), BUF_SIZE,ifh)) > 0)
	{
		dbg("Read %d bytes, passing through CipherUpdate...\n", read_size);
		if(EVP_CipherUpdate(ctx, outbuf, &len, inbuf, read_size) == 0) {
			rc = AES_ERR_CIPHER_UPDATE;
			fprintf(stderr, "ERROR: EVP_CipherUpdate failed. OpenSSL error: %s\n", ERR_error_string(ERR_get_error(), NULL));
			goto cleanup;
		}
		dbg("\tGot back %d bytes from CipherUpdate...\n", len);
		
		dbg("Writing %d bytes to %s...\n", len, outfile);
		if(_write(outbuf,sizeof(char),len,ofh) != len) {
			rc = AES_ERR_IO;
			fprintf(stderr, "ERROR: Writing to the file %s failed. errno = %s\n", outfile, strerror(errno));
			goto cleanup;
		}
		dbg("\tWrote %d bytes\n", len);
		
		u_len += len;
	}
	
	// Check last read succeeded
	if(read_size == -1) {
		rc = AES_ERR_IO;
		fprintf(stderr, "ERROR: Reading from the file %s failed. errno = %s\n", infile, strerror(errno));
		goto cleanup;
	}
	
	// Finalize encryption/decryption
	rc = AES_ERR_CIPHER_FINAL;
	if(EVP_CipherFinal_ex(ctx, outbuf, &f_len) == 0) {
		fprintf(stderr, "ERROR: EVP_CipherFinal_ex failed. OpenSSL error: %s\n", ERR_error_string(ERR_get_error(), NULL));
		goto cleanup;
	}
	
	dbg("u_len = %d, f_len = %d\n", u_len, f_len);
	
	// Write the final block, if any
	if(f_len) {
		dbg("Writing final %d bytes to %s...\n", f_len, outfile);
		if(_write(outbuf,sizeof(char),f_len,ofh) != f_len) {
			rc = AES_ERR_IO;
			fprintf(stderr, "ERROR: Final write to the file %s failed. errno = %s\n", outfile, strerror(errno));
			goto cleanup;
		}
		dbg("\tWrote last %d bytes\n", f_len);
	}

	rc = u_len + f_len;

 cleanup:
	EVP_CIPHER_CTX_free(ctx);
 	//EVP_CIPHER_CTX_cleanup(ctx);
 	if(ifh != NULL){
	 	if (0==inmod) 
		 	pclose(ifh); 
		else
			fclose(ifh);
	}

 	if(ofh != NULL){
	 	if (0==outmod){
			pclose(ofh);
		 }
		else if (1==outmod)
		{
			fclose(ofh);
		}
	}
 	
	return rc;
}


static PyObject* cry_cryptofile(PyObject *dummy, PyObject* args) {
	cryptomaniac_t cm;
	int st;
	char * p1,*p2;
	unsigned char *p3=NULL;
	int rst;
	char rsts[256];
	PyObject *result;
	rst=1;
	if (PyArg_ParseTuple(args,"ss|s",&p1,&p2,&p3)){
		cm.infile=p1;
		cm.outfile=p2;
		getIkeyIv(p3,cm.key,cm.iv);
		cm.encrypt=1;
		cm.mode=EVP_get_cipherbyname(EVP_CIPHERNAME_AES_CTR);

		if((st = aes_encrypt_file(cm.infile, cm.outfile, cm.key, cm.iv, cm.mode, cm.encrypt,1,1)) <= 0)
		{
			sprintf(rsts, "ERROR: %s failed\n", cm.encrypt ? "Encryption" : "Decryption");
			rst=0;
		}else
			sprintf(rsts,"OK");
	}
 	result=Py_BuildValue("(i,s)",rst,rsts);
	return result;
}
static PyObject* cry_zip(PyObject *dummy,PyObject* args){
cryptomaniac_t cm;
	int st;
	char * p1,*p2;
	unsigned char *p3=NULL;
	int rst;
	char rsts[256];
	PyObject *result;
	rst=1;
	if (PyArg_ParseTuple(args,"ss|s",&p1,&p2,&p3)){
		sprintf(rsts,"tar cj -C %s .",p1);
		cm.infile=rsts;
		cm.outfile=p2;
		getIkeyIv(p3,cm.key,cm.iv);
		cm.encrypt=1;
		cm.mode=EVP_get_cipherbyname(EVP_CIPHERNAME_AES_CTR);

		if((st = aes_encrypt_file(cm.infile, cm.outfile, cm.key, cm.iv, cm.mode, cm.encrypt,0,1)) <= 0)
		{
			sprintf(rsts, "ERROR: %s failed\n", cm.encrypt ? "Encryption" : "Decryption");
			rst=0;
		}else
			sprintf(rsts,"OK");
	}
 	result=Py_BuildValue("(i,s)",rst,rsts);
	return result;
}

static PyObject* cry_unzip(PyObject *dummy,PyObject* args){
	cryptomaniac_t cm;
	int st;
	char * p1,*p2;
	unsigned char *p3=NULL;
	int rst;
	char rsts[256];
	PyObject *result;
	rst=1;
	if (PyArg_ParseTuple(args,"ss|s",&p1,&p2,&p3)){
		sprintf(rsts,"tar xj -C %s",p2);
		cm.infile=p1;
		cm.outfile=rsts;
		getIkeyIv(p3,cm.key,cm.iv);
		cm.encrypt=0;
		cm.mode=EVP_get_cipherbyname(EVP_CIPHERNAME_AES_CTR);

		if((st = aes_encrypt_file(cm.infile, cm.outfile, cm.key, cm.iv, cm.mode, cm.encrypt,1,0)) <= 0)
		{
			sprintf(rsts, "ERROR: %s failed\n", cm.encrypt ? "Encryption" : "Decryption");
			rst=0;
		}else
			sprintf(rsts,"OK");
	}
 	result=Py_BuildValue("(i,s)",rst,rsts);
	return result;
}
static PyObject* cry_getMD5(PyObject *dummy,PyObject* args){
	unsigned char * p1;
	unsigned char buf[16];
	char md5[32];
	PyObject *result;
	int i ;
	if (PyArg_ParseTuple(args,"s",&p1)){
		MD5(p1,strlen((char*)p1),buf);
		for (i=0;i<16;i++)
			sprintf(&md5[i<<1],"%02x",buf[i]);
		result=Py_BuildValue("s",md5);
		return result;
	}
 	result=Py_BuildValue("s","");
	return result;
}
static PyObject* cry_callback(PyObject *dummy, PyObject *args)
{
	cryptomaniac_t cm;
	int st;
	PyObject * callback;
	char * p1;
	unsigned char *p3=NULL;
	char buf[256];
	PyObject *result;
	PyObject *arglist;
	if (PyArg_ParseTuple(args,"Os|s:callbacka",&callback,&p1,&p3)){
		if (!PyCallable_Check(callback)) {
            PyErr_SetString(PyExc_TypeError, "parameter must be callable");
            return NULL;
        }
		cm.infile=p1;
		getIkeyIv(p3,cm.key,cm.iv);
		cm.encrypt=0;
		cm.mode=EVP_get_cipherbyname(EVP_CIPHERNAME_AES_CTR);

		if((st = aes_encrypt_file(cm.infile, cm.outfile, cm.key, cm.iv, cm.mode, cm.encrypt,1,2)) <= 0)
		{
			sprintf(buf, "ERROR: %s failed\n", cm.encrypt ? "Encryption" : "Decryption");
			PyErr_SetString(PyExc_BaseException, buf);
			mwrite(NULL,0,0,(FILE*)3);
			return NULL;
		}else{
			char *c;
			size_t len;
			len=mwrite(&c,0,0,(FILE*)1);
			arglist = Py_BuildValue("(y#)", c,len);
			result = PyObject_CallObject(callback, arglist);
			Py_DECREF(arglist);
			mwrite(NULL,0,0,(FILE*)3);
			return result;
		}
	}else
	{
		PyErr_SetString(PyExc_Exception, "parameter analysis error.");
        return NULL;
	}
}

static struct PyMethodDef cry_methods[] = {
    {"cryptofile",(PyCFunction)cry_cryptofile,METH_VARARGS,NULL},
	{"zip",(PyCFunction)cry_zip,METH_VARARGS,NULL},
	{"unzip",(PyCFunction)cry_unzip,METH_VARARGS,NULL},
	{"getMD5",(PyCFunction)cry_getMD5,METH_VARARGS,NULL},
    {"callback",(PyCFunction)cry_callback,METH_VARARGS,NULL},
    {NULL,NULL,0,NULL}
};

static struct PyModuleDef spammodule = {
    PyModuleDef_HEAD_INIT,
    "cry",   /* name of module */
    NULL, /* module documentation, may be NULL */
    -1,       /* size of per-interpreter state of the module,
                 or -1 if the module keeps state in global variables. */
    cry_methods
};

PyMODINIT_FUNC PyInit_cry(void)
{
    return PyModule_Create(&spammodule);
}
